#include <stdlib.h>
#include <stdio.h>
#include <Windows.h>
#include <sys/time.h>

#define a(i,j) a[(i)*k + (j)]
#define b(i,j) b[(i)*n + (j)]
#define c(i,j) c[(i)*n + (j)]

// 4x4分块计算，使用寄存器
void mul_4_4_ker(float *pakage_a, float *pakage_b, float *c, int x, int y, int k, int n){
    register float c00, c01, c02, c03,
                c10, c11, c12, c13,
                c20, c21, c22, c23,
                c30, c31, c32, c33;

    c00 = c01 = c02 = c03 = 0.0;
    c10 = c11 = c12 = c13 = 0.0;
    c20 = c21 = c22 = c23 = 0.0;
    c30 = c31 = c32 = c33 = 0.0;

    for (int i = 0; i < k; i++){
        int a1 = i + 1*k;
        int a2 = i + 2*k;
        int a3 = i + 3*k;
        
        c00 += pakage_a[i] * pakage_b[i];
        c01 += pakage_a[i] * pakage_b[a1];
        c02 += pakage_a[i] * pakage_b[a2];
        c03 += pakage_a[i] * pakage_b[a3];

        c10 += pakage_a[a1] * pakage_b[i];
        c11 += pakage_a[a1] * pakage_b[a1];
        c12 += pakage_a[a1] * pakage_b[a2];
        c13 += pakage_a[a1] * pakage_b[a3];
        
        c20 += pakage_a[a2] * pakage_b[i];
        c21 += pakage_a[a2] * pakage_b[a1];
        c22 += pakage_a[a2] * pakage_b[a2];
        c23 += pakage_a[a2] * pakage_b[a3];
        
        c30 += pakage_a[a3] * pakage_b[i];
        c31 += pakage_a[a3] * pakage_b[a1];
        c32 += pakage_a[a3] * pakage_b[a2];
        c33 += pakage_a[a3] * pakage_b[a3];
    }

    c(x,y) += c00;
    c(x,y+1) += c01;
    c(x,y+2) += c02;
    c(x,y+3) += c03;

    c(x+1,y) += c10;
    c(x+1,y+1) += c11;
    c(x+1,y+2) += c12;
    c(x+1,y+3) += c13;

    c(x+2,y) += c20;
    c(x+2,y+1) += c21;
    c(x+2,y+2) += c22;
    c(x+2,y+3) += c23;

    c(x+3,y) += c30;
    c(x+3,y+1) += c31;
    c(x+3,y+2) += c32;
    c(x+3,y+3) += c33;
}

void pack_A(float *a, int k, int m, float *pakage_a){
    float *tep = &a(m,0);
    for (int i = 0; i < k; i++){
        pakage_a[i] = tep[i];
        pakage_a[i+k] = tep[i+k];
        pakage_a[i+2*k] = tep[i+2*k];
        pakage_a[i+3*k] = tep[i+3*k];
    }
}

void pack_B(float *b, int k, int n, int col, float *pakage_b){
   for (int i = 0; i < k; i++){
        int s = i*n + col;
        pakage_b[i] = b[s];
        pakage_b[i+k] = b[s +1];
        pakage_b[i+2*k] = b[s +2];
        pakage_b[i+3*k] = b[s +3];

   }
}

void matmul_4_4(int m, int n, int k, float *a, float *b, float *c){
    float *p_a = (float *)malloc(4*k*sizeof(float));
    float *p_b = (float *)malloc(4*k*sizeof(float));
    for (int i = 0; i < m; i+=4){
        pack_A(a, k, i, p_a);
        for (int j = 0; j < n; j+=4){
            pack_B(b, k, n, j, p_b);
            mul_4_4_ker(p_a, p_b, c, i, j, k, n);
        }
    }
    
    free(p_a);
    free(p_b);
}


typedef struct Mes{
    float *a;
    float *b;
    float *c;
    int m;
    int n;
    int k;
}*mes_ptr;

// 多线程处理
// windows

DWORD WINAPI mul_4_4_thread_function(LPVOID lpParam){
    mes_ptr p = (struct Mes *)lpParam;
    matmul_4_4(p -> m, p -> n, p -> k, p -> a, p -> b, p -> c);
    return 0;
}

void mul_4_4_mp(int m, int n, int k, float *a, float *b, float *c){
    const int num_threads = 16;
    struct Mes mes[num_threads];
    HANDLE h[num_threads];
    int m_nums = m/num_threads;
    
    for (int i = 0; i < num_threads; i++){
        mes[i].a = &a(i*m_nums,0);
        mes[i].b = b;
        mes[i].c = &c(i*m_nums,0);
        mes[i].m = m_nums;
        mes[i].n = n;
        mes[i].k = k;
    }

    for (int i = 0; i < num_threads; i++){
        h[i] = CreateThread(NULL, 0, mul_4_4_thread_function, &mes[i], 0, NULL);
    }
    
    WaitForMultipleObjects(num_threads, h, TRUE, INFINITE);

    for (int i = 0; i < num_threads; i++){
        CloseHandle(h[i]);
    }

}

// linux
// #include <pthread.h>
// void *mul_4_4_thread_function(void *p){
//     struct Mes *mes = (struct Mes *)p;
//     matmul_4_4(mes -> m, mes -> n, mes -> k, mes -> a, mes -> b, mes -> c);
// }

// void mul_4_4_mp(int m, int n, int k, float *a, float *b, float *c){
//     const int num_threads = 8;
//     struct Mes mes[num_threads];
//     pthread_t threads[num_threads];
//     int m_nums = m/num_threads;
    
//     for (int i = 0; i < num_threads; i++){
//         mes[i].a = &a(i,0);
//         mes[i].b = b;
//         mes[i].c = &c(i,0);
//         mes[i].m = m_nums;
//         mes[i].n = n;
//         mes[i].k = k;
//     }

//     for(int i = 0; i < num_threads; i++){
//         pthread_create(&threads[i], NULL, mul_4_4_thread_function, &mes[i]);
//     }

//     for (int i = 0; i < num_threads; i++) {
//         pthread_join(threads[i], NULL);
//     }

// }


int main(){
    int m = 2048, n = 2048, k = 2048;
    float *a = (float *)malloc(m*k*sizeof(float));
    float *b = (float *)malloc(k*n*sizeof(float));
    float *c = (float *)malloc(m*n*sizeof(float));

    for (int i = 0; i < m*k; i++){
        a[i] = 1.0f;
    }

    for (int i = 0; i < k*n; i++){
        b[i] = 1.0f;
    }

    for (int i = 0; i < m*n; i++){
        c[i] = 0.0f;
    }

    // struct timeval start, end;
    // long seconds_finished;
    // unsigned long micro_seconds_finished;
    // double elapsed;
    // gettimeofday(&start, NULL);
    // gettimeofday(&end, NULL);
    // seconds_finished = end.tv_sec - start.tv_sec;
    // micro_seconds_finished = end.tv_usec - start.tv_usec;
    // elapsed = ((seconds_finished * 1000000) + micro_seconds_finished) / 1000000.0;
    // printf("time: %f s\n", elapsed);


    DWORD start_t, end_t;

    start_t = GetTickCount();
    
    mul_4_4_mp(m , n, k, a, b, c);

    end_t = GetTickCount();

    double time = (end_t - start_t)/1e+3;
    printf("time: %lf s\n", time);

    double flops = 2.0 * m * n * k * 1e-9;
    printf("flops: %lf Gflops\n", flops/time);

    // for (int i = 0; i < m; i++){
    //     for (int j = 0; j < n; j++){
    //         printf("%lf ", c(i,j));
    //         }
    //     printf("\n");
    // }
    free(a);
    free(b);
    free(c);

    return 0;
}